package fun.yyds.config.websocket;

import com.fasterxml.jackson.databind.ObjectMapper;
import org.springframework.stereotype.Component;
import org.springframework.web.socket.TextMessage;
import org.springframework.web.socket.WebSocketSession;
import org.springframework.web.socket.handler.TextWebSocketHandler;

import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.concurrent.*;
import java.util.concurrent.locks.ReentrantLock;

@Component
public class DataPushHandler extends TextWebSocketHandler {
    // 存储用户会话和对应的推送任务信息
    private final Map<String, SessionInfo> sessionInfoMap = new ConcurrentHashMap<>();
    private final ObjectMapper objectMapper = new ObjectMapper();
    private final ReentrantLock lock = new ReentrantLock();

    // 会话信息类，保存推送状态、任务等
    private static class SessionInfo {
        WebSocketSession session;
        ScheduledFuture<?> pushTask;
        List<Data> dataList;
        int currentIndex;
        boolean isPaused;
        boolean hasMoreData;  // 是否还有更多数据需要查询
        QueryParams queryParams;  // 保存查询参数，用于继续查询
        int currentPage;  // 当前查询页码

        SessionInfo(WebSocketSession session) {
            this.session = session;
            this.dataList = new ArrayList<>();
            this.currentIndex = 0;
            this.isPaused = false;
            this.hasMoreData = true;
            this.currentPage = 0;
        }
    }

    // 连接建立时触发
    @Override
    public void afterConnectionEstablished(WebSocketSession session) {
        String sessionId = session.getId();
        sessionInfoMap.put(sessionId, new SessionInfo(session));
        System.out.println("WebSocket连接建立，会话ID：" + sessionId);
    }

    // 开始推送数据
    public void startPush(String sessionId, QueryParams params, List<Data> initialData) {
        SessionInfo sessionInfo = sessionInfoMap.get(sessionId);
        if (sessionInfo == null) {
            return;
        }

        // 初始化会话信息
        sessionInfo.queryParams = params;
        sessionInfo.dataList = initialData;
        sessionInfo.currentIndex = 0;
        sessionInfo.isPaused = false;
        sessionInfo.currentPage = 0;
        sessionInfo.hasMoreData = initialData.size() >= 10000;  // 如果首次查询返回10000条，说明可能还有更多数据

        // 取消之前的任务（如果存在）
        if (sessionInfo.pushTask != null && !sessionInfo.pushTask.isCancelled()) {
            sessionInfo.pushTask.cancel(false);
        }

        startPushTask(sessionInfo);
    }

    // 暂停推送
    public void pausePush(String sessionId) {
        SessionInfo sessionInfo = sessionInfoMap.get(sessionId);
        if (sessionInfo != null) {
            sessionInfo.isPaused = true;
            System.out.println("会话 " + sessionId + " 已暂停推送");
        }
    }

    // 继续推送
    public void resumePush(String sessionId) {
        SessionInfo sessionInfo = sessionInfoMap.get(sessionId);
        if (sessionInfo != null && sessionInfo.isPaused) {
            sessionInfo.isPaused = false;
            System.out.println("会话 " + sessionId + " 已恢复推送");
        }
    }

    // 启动推送任务
    private void startPushTask(SessionInfo sessionInfo) {
        ScheduledExecutorService executor = Executors.newSingleThreadScheduledExecutor();

        sessionInfo.pushTask = executor.scheduleAtFixedRate(() -> {
            try {
                // 如果暂停状态，不执行推送
                if (sessionInfo.isPaused) {
                    return;
                }

                WebSocketSession session = sessionInfo.session;
                if (session == null || !session.isOpen()) {
                    cleanupSession(session.getId());
                    return;
                }

                // 检查是否需要加载更多数据
                if (sessionInfo.currentIndex >= sessionInfo.dataList.size() && sessionInfo.hasMoreData) {
                    loadMoreData(sessionInfo);
                }

                // 如果数据已推送完毕，关闭任务和连接
                if (sessionInfo.currentIndex >= sessionInfo.dataList.size() && !sessionInfo.hasMoreData) {
                    sessionInfo.pushTask.cancel(false);
                    executor.shutdown();
                    session.close();
                    cleanupSession(session.getId());
                    return;
                }

                // 推送当前批次数据（最多5条）
                int endIndex = Math.min(sessionInfo.currentIndex + 5, sessionInfo.dataList.size());
                List<Data> batch = sessionInfo.dataList.subList(sessionInfo.currentIndex, endIndex);
                sessionInfo.currentIndex = endIndex;

                // 发送数据
                String jsonData = objectMapper.writeValueAsString(batch);
                session.sendMessage(new TextMessage(jsonData));

            } catch (Exception e) {
                System.err.println("推送失败：" + e.getMessage());
                sessionInfo.pushTask.cancel(false);
                executor.shutdown();
                cleanupSession(sessionInfo.session.getId());
            }
        }, 0, 5000, TimeUnit.MILLISECONDS);  // 每5s执行一次
    }

    // 加载更多数据（分页查询）
    private void loadMoreData(SessionInfo sessionInfo) {
        try {
            lock.lock();
            // 防止并发加载
            if (sessionInfo.currentIndex < sessionInfo.dataList.size()) {
                return;
            }

            sessionInfo.currentPage++;
            // 这里实际应用中应该调用数据服务查询下一页数据
            // 为了演示，这里只是模拟
            List<Data> moreData = new ArrayList<>();
            // moreData = dataService.queryMoreData(sessionInfo.queryParams, sessionInfo.currentPage, 10000);

            // 如果返回的数据少于10000条，说明没有更多数据了
            sessionInfo.hasMoreData = moreData.size() >= 10000;
            sessionInfo.dataList.addAll(moreData);
            System.out.println("已加载第 " + sessionInfo.currentPage + " 页数据，共 " + moreData.size() + " 条");

        } finally {
            lock.unlock();
        }
    }

    // 清理会话资源
    private void cleanupSession(String sessionId) {
        SessionInfo sessionInfo = sessionInfoMap.remove(sessionId);
        if (sessionInfo != null && sessionInfo.pushTask != null) {
            sessionInfo.pushTask.cancel(false);
        }
        System.out.println("会话 " + sessionId + " 已清理");
    }

    // 连接关闭时清理资源
    @Override
    public void afterConnectionClosed(WebSocketSession session, org.springframework.web.socket.CloseStatus status) {
        cleanupSession(session.getId());
        System.out.println("WebSocket连接关闭，会话ID：" + session.getId());
    }

    // 获取会话ID
    public String getSessionId(WebSocketSession session) {
        return session.getId();
    }
}
